import os

from benchmarks.maxcut import Benchmark
import numpy as np
import matplotlib.pyplot as plt
import multiprocessing as mp

from utils.md5 import md5
from utils.tsne import tsne_plot
from optimizer.qsbpso import Optimizer as QsbpsoOptimizer
from optimizer.sb import Optimizer as SbOptimizer
from tqdm import tqdm
import pickle
import pandas as pd

config = {
    'n_dim': 2000
}
bm = Benchmark(config)

config = {
    'n_run': 6000,
    'n_part': 10,
    'show': 0,
    'benchmark': bm,
    'n_group': 1,
    'n_dim': bm.ising_equivalent.dimension,
    'pos_max': 2,
    'pos_min': -2,
    'max_v': 2,
    'min_v': -2,
    'fe_max': 60000,
    'record_per_fe': 300,
    'ising_J': bm.ising_equivalent.J,
}


def evaluate(Optimizer):
    ress = []
    for _ in tqdm(range(20)):
        optimizer = Optimizer(config)
        optimizer.run()
        print(optimizer.result_cache)
        best_res = optimizer.result_cache[-1]['best']
        ress.append(best_res)
    return {'name': Optimizer.optimizer_name, 'res': {
        'mean': np.mean(ress),
        'best': np.min(ress),
        'all': ress,
        'std': np.std(ress)
    }}


def evaluate_with_config(config):
    m = md5(str(config))
    result_path = f'data/run_result/{m}.pickle'
    if os.path.exists(result_path):
        # logger.info(f'{task_md5} cache yes')
        with open(result_path, 'rb') as f:
            return pickle.load(f)
    Optimizer = config['Optimizer']
    config_detail = {
        'n_run': 6000,
        'n_part': 10,
        'show': 0,
        'benchmark': bm,
        'n_group': 1,
        'n_dim': bm.ising_equivalent.dimension,
        'pos_max': 2,
        'pos_min': -2,
        'max_v': 2,
        'min_v': -2,
        'fe_max': 60000,
        'record_per_fe': 300,
        'ising_J': bm.ising_equivalent.J,
        **config
    }

    ress = []
    for _ in tqdm(range(10)):
        optimizer = Optimizer(config_detail)
        optimizer.run()
        print(optimizer.result_cache)
        best_res = optimizer.result_cache[-1]['best']
        ress.append(best_res)

    result = {'name': Optimizer.optimizer_name,
              'mean': np.mean(ress),
              'best': np.min(ress),
              # 'all': ress,
              'std': np.std(ress),
              **config
              }
    with open(result_path, 'wb') as f:
        pickle.dump(result, f, 0)
    return result


if __name__ == '__main__':
    configs = []
    base_config = {
        'Optimizer': QsbpsoOptimizer,
        'crossover_rate': 0,
        'mutation_rate': 0,
        'sa_rate': 0,
        'fine_tune_rate': 0,
        'pso_rate': 0,
        'no_improve_limit': 150,
        'n_part': 10,
        'n_run': 1e9,
    }
    configs.append({'Optimizer': SbOptimizer, 'n_part': 1, 'n_run': 1e9})
    configs.append({'Optimizer': SbOptimizer, 'n_part': 10, 'n_run': 1e9})
    configs.append({'Optimizer': SbOptimizer, 'n_part': 100, 'n_run': 1e9})
    # self.crossover_rate = 0.001
    # self.mutation_rate = 0.001
    # self.sa_rate = 0.001
    # self.fine_tune_rate = 0.001
    # self.pso_rate = 0.001
    # self.T = 100
    # self.no_improve_limit = 100
    for crossover_rate in [0, 1e-2, 1e-4]:
        configs.append({
            'Optimizer': QsbpsoOptimizer,
            'crossover_rate': crossover_rate,
        })

    for sa_rate in [0, 1e-2, 1e-4]:
        configs.append({
            'Optimizer': QsbpsoOptimizer,
            'sa_rate': sa_rate,
        })
    for fine_tune_rate in [0, 1e-2, 1e-4]:
        configs.append({
            'Optimizer': QsbpsoOptimizer,
            'fine_tune_rate': fine_tune_rate,
        })
    for pso_rate in [0, 1e-2, 1e-4]:
        configs.append({
            'Optimizer': QsbpsoOptimizer,
            'pso_rate': pso_rate,
        })
    for no_improve_limit in [1, 10, 100, 1000]:
        configs.append({
            'Optimizer': QsbpsoOptimizer,
            'no_improve_limit': no_improve_limit,
            'crossover_rate': 1e-5,
            'sa_rate': 1e-5,
            'fine_tune_rate': 1e-5,
            'pso_rate': 1e-5,
        })
    for n_part in [10, 30, 100]:
        configs.append({
            'Optimizer': QsbpsoOptimizer,
            'crossover_rate': 1e-5,
            'sa_rate': 1e-5,
            'fine_tune_rate': 1e-5,
            'pso_rate': 1e-5,
            'n_part': n_part,
        })
    for mutation_rate in [0, 1e-2, 1e-4]:
        configs.append({
            'Optimizer': QsbpsoOptimizer,
            'mutation_rate': mutation_rate,
        })
    print(f'任务总数:{len(configs)}')
    with mp.Pool(processes=3) as pool:  # start 4 worker processes
        res = pool.map(evaluate_with_config, configs)  # prints "[0, 1, 4,..., 81]"
        print(res)
        df = pd.DataFrame(res)
        # print(df)
        writer = pd.ExcelWriter('最终结果.xlsx')
        df.to_excel(writer, 'Sheet1')  # 这里假设df是一个pandas的dataframe
        writer.save()
        writer.close()
